import os
import argparse
from Grade import DeeplabGrade
import cv2
import time

model = DeeplabGrade("model_241.pth")


def run_one(file):
    t1 = time.perf_counter()
    mask_img, prediction = model.infer(file, class_colors)
    t2 = time.perf_counter()
    print("Infer file: {}, cost: {}s".format(os.path.basename(file), (t2 - t1)))
    return mask_img, prediction


def loop(input: str, output: str):
    for root, dirs, files in os.walk(input):
        for dir_name in dirs:
            predict_path = os.path.join(root, dir_name)
            output_folder = os.path.join(output, dir_name, 'pred')
            output_folder_cal = os.path.join(output, dir_name, 'cal')

            os.makedirs(output_folder, exist_ok=True)
            os.makedirs(output_folder_cal, exist_ok=True)

            for root1, dirs1, files1 in os.walk(predict_path):
                for predict_img in files1:
                    text, extension = os.path.splitext(predict_img)
                    img_path = os.path.join(root1, predict_img)
                    colored_mask_img, predictions = run_one(img_path)
                    result_path = os.path.join(output_folder, f"{text}.png")
                    colored_mask_img.save(result_path, format='PNG')
                    result_path_cal = os.path.join(output_folder_cal, f"{text}.png")
                    cv2.imwrite(result_path_cal, predictions)
                    print(f"Saved result for {predict_img} at {result_path}")


def run(workdir: str):
    input = os.path.join(workdir, "user_input")
    output = os.path.join(workdir, "user_output")
    if not os.path.exists(output):
        os.makedirs(output)
    output = os.path.join(workdir, "user_output")
    if not os.path.exists(input):
        print("user_input not found")
    loop(input, output)


if __name__ == "__main__":
    class_colors = {
        1: (255, 0, 0),  # Red
        2: (0, 255, 0),  # Green
        3: (0, 0, 255),  # Blue
        4: (255, 255, 0),  # Yellow
        5: (255, 0, 255)  # Magenta
    }
    parser = argparse.ArgumentParser()
    parser.add_argument('--workdir', type=str, required=True)
    args = parser.parse_args()
    if args.workdir == '':
        print("workdir is empty")

    run(args.workdir)
